from Tkinter import *
from os import getcwd
from tkFileDialog import askopenfilename
import numpy as np
import numpy.matlib
import math
import tkMessageBox

# ------ interface ------

class Parcourir(Frame):
    'Search for the file and get the filename'
    def __init__(self, master = None, command = None, **kw):
        Frame.__init__(self, master, **kw)

        self.chemin = StringVar()
        self.chemin.set(getcwd())    
        
        self.command = command     

        self.entree_chemin = Entry(self, text = self.chemin,bg='lightgreen',fg='blue',relief='sunken',)
        self.entree_chemin.pack(side='top',fill=X)

        self.bouton = Button(self,text = "Parcourir",bg='lightblue',relief='raised',border=2,command = self.parcours)
        self.bouton.pack(side='right')

    def parcours(self):
        chemin = askopenfilename()
        if chemin != '':
            self.chemin.set(chemin)
            if self.command != None:
                self.command()
                
    def get(self):
        return str(self.chemin.get())

        
def demo_parcours():
    global label_parcours, parcours
    
    fen_parcours = Toplevel(fen)
    fen_parcours.grab_set()

    lab1=Label(fen_parcours,relief='groove',border=2,text='Address of the file *.cube is:',width=80)
    lab1.pack(side='top')

    label_parcours = Label(fen_parcours)
    label_parcours.pack(side='top',expand='true')

    parcours = Parcourir(fen_parcours, command\
                         = modify_label)
    parcours.pack(side='top',fill=X)
    modify_label()


def modify_label():
    label_parcours.config(text = parcours.get())
    #parcours.get():the filename is stocked in this global variable

# ------ Algorithme functions ------


# --- 1. basic information ---

#Principal fonction in this part,
#and this is the command for the button "Information of the molecule"
def moleinfo():
    """get all the info of the molecule"""
    global mole
    initi()
    mole=Molef()
    

def initi():
   """read the title of the document .cube"""
   global fI
   fI=open(parcours.get(),'r') 
   fI.readline()
   fI.readline()

class PosAtM():
    "bind the number of an atom to its position in a molecule????"
    def __init__(self):
        self.nb=0 #number of the atom
        self.ch=0 #charge of the atom
        self.x=0
        self.y=0
        self.z=0
    def affiche(self):
        print '[',self.nb,self.ch,self.x,self.y,self.z,']'
    
class Molef():
    "all the information of the molecule"
    def __init__(self):
        self.temp=self.lignetraite()  
        self.NbTa=int(self.temp[0]) #total number of atoms
        
        #the position of the origin
        self.ori=[0]*3
        for i in range(3): 
            self.ori[i]=float(self.temp[i+1])

        #step infomation in three directions
        self.nbpas=[0]*3
        self.pas=[0]*3
        for i in range(3):
            self.temp=self.lignetraite()
            self.nbpas[i]=int(self.temp[0])   #the number of steps in each direction
            self.pas[i]=float(self.temp[i+1]) #the distance of step in each direction

        #information of the atoms in the molecule
        self.atom=[0]*self.NbTa      #create a list to stock the atom information
        for i in range(self.NbTa):
            self.atom[i]=PosAtM()
            self.temp=self.lignetraite()
            self.atom[i].nb=int(self.temp[0])
            self.atom[i].ch=float(self.temp[1])
            self.atom[i].x=float(self.temp[2])
            self.atom[i].y=float(self.temp[3])
            self.atom[i].z=float(self.temp[4])

    def affiche(self):
        for i in range(self.NbTa):
            self.atom[i].affiche()
            
    def lignetraite(self): #treatment of a line
        v=fI.readline()
        v=v.split()     
        return v

def sautpart1():
    fI=open(parcours.get(),'r')
    n=6+mole.NbTa   
    for i in range(n):
        fI.readline()
    

# ---- 1. End of the reading of basic information(tested)----

# --- 2. search for the value of a point in the .cube ---

## --2.1 Arrange all the numbers in the order that it appeared in the .cube --???
##Quicksort
from random import *

def partition(a, start, end, pivotIndex):
    low = start
    high = end - 1  # After we remove pivot it will be one smaller
    pivotValue = a[pivotIndex].vlsort
    temp=a[pivotIndex]
    a[pivotIndex] = a[end]
    while True:
        while low <= high and a[low].vlsort < pivotValue:
            low = low + 1
        while low <= high and a[high].vlsort >= pivotValue:
            high = high - 1
        if low > high:
            break
        a[low], a[high] = a[high], a[low]
    a[end] = a[low]
    a[low] = temp
    return low



def insertionSort(a, start, end):
    for i in xrange(start, end + 1):
        # Insert a[i] into the sorted sublist
        v = a[i]
        for j in reversed(xrange(0, i)):
            if a[j].vlsort <= v.vlsort:
                a[j + 1] = v
                break
            a[j + 1] = a[j]
        else:
            a[0] = v
    return a


def qsortRange(a, start, end):
    if end - start + 1 < 32:
        insertionSort(a, start, end)
    else:
        pivotIndex = partition(a, start, end, randint(start, end))
        qsortRange(a, start, pivotIndex - 1)
        qsortRange(a, pivotIndex + 1, end)
    return a

def calvlsort1(a):
    for i in range(len(a)):
        a[i].vlsort=a[i].pos1[2]+a[i].pos1[1]*mole.nbpas[2]+\
                     a[i].pos1[0]*mole.nbpas[2]*mole.nbpas[1]

def calvlsort2(a):
    for i in range(len(a)):
        a[i].vlsort=a[i].pos2[2]+a[i].pos2[1]*mole.nbpas[2]+\
                     a[i].pos2[0]*mole.nbpas[2]*mole.nbpas[1]

def qsort(a,k): #k is the parameter to choose the corresponding serie of positions
    if k==1:
        calvlsort1(a)
    if k==2:
        calvlsort2(a)
    return qsortRange(a, 0, len(a) - 1)

## -- 2.1 all the numbers are in the order that it appeared in the .cube (test)--


## -- 2.2 search the value of these numbers in the .cube --

def Ptcube(a): #a=Ptvl in this section
    """this function is to locate all the points concerned in the .cube"""
    for i in range(len(a)):
        for j in drc:
            a[i].pos1[j]=a[i].pos[j]
            a[i].pos2[j]=a[i].pos[j]
        a[i].pos1[flag]=int(a[i].pos[flag])
        a[i].pos2[flag]=int(a[i].pos[flag])+1
        
        a[i].r1=a[i].pos[flag]-a[i].pos1[flag]
        a[i].r2=1-a[i].r1

def nblignebloc(m):
    """the number of lines in a bloc, i.e. a group of m values"""
    nbligne=int(m/6)
    if m%6 is not 0:
        nbligne=nbligne+1
    return nbligne

def getval1(diff,onvalue):
    """the most simple case to get a value"""
    global vline
    
    if onvalue==0:
        vline=fI.readline()
        vline=vline.split()
        temp=float(vline[diff[0]-1])
        return temp
    
    for i in range(diff[1]-1):
        fI.readline()
    if diff[0]==0:
        v=fI.readline()
        v=v.split()
        return float(v[5])
    elif diff[0] is not 0:
        vline=fI.readline()
        vline=vline.split()
        temp=float(vline[diff[0]-1])
        return temp
        

def getval(pospre,poscur):
    """get the value at the current postion 'poscur'
       from the precedent position 'pospre'"""
    diff=[0]*2
    #0<=diff[0]<=6, which is the horizontal translation from a number to another
    #diff[1] is the number of lines to pass
    
    if  pospre[2]==mole.nbpas[2]-1:
        diff[0]=(poscur[2]+1)%6
        diff[1]=(poscur[2]+1)/6+(poscur[1]-pospre[1]-1)*nbligneb+\
                 (poscur[0]-pospre[0])*mole.nbpas[1]*nbligneb
        val=getval1(diff,1)
        return val

    elif (pospre[2]+1)%6==0:
        diff[0]=(poscur[2]+1)%6
        diff[1]=(poscur[2]+1)/6-(pospre[2]+1)/6+(poscur[1]-pospre[1])*nbligneb+\
                 (poscur[0]-pospre[0])*mole.nbpas[1]*nbligneb
        val=getval1(diff,1)
        return val

    elif pospre[0]==poscur[0] and pospre[1]==poscur[1]:
        if pospre[2]/6==poscur[2]/6: #in the same line
            CurIndex=(poscur[2]+1)%6
            if CurIndex==0:
                CurIndex=6
            try:
                val=float(vline[CurIndex-1])
            except:
                diff=[CurIndex,0]
                val=getval1(diff,0)
        else:
            diff[1]=poscur[2]/6-pospre[2]/6-1
            diff[0]=(poscur[2]+1)%6
            val=getval1(diff,1)
            
        return val

    else:
        diff[1]=((poscur[2]+1)/6-(pospre[2]+1)/6-1)+(poscur[1]-pospre[1])*nbligneb+\
                 (poscur[0]-pospre[0])*mole.nbpas[1]*nbligneb
        diff[0]=(poscur[2]+1)%6
        val=getval1(diff,1)
        return val
       
def searchval1(a):
    """This function is to search for the value on the position 1"""
    qsort(a,1)

    global nbligneb
        
    n=0
    nbligneb=nblignebloc(mole.nbpas[2])
    pospre=[0,0,0] #initialise the precedent postion
    for i in range(mole.nbpas[0]):
        if n>=len(a):
            return
        if i==a[n].pos1[0]:
            for j in range(mole.nbpas[1]):
                if n>=len(a):
                    return
                if j==a[n].pos1[1]:
                    for k in range(mole.nbpas[2]):
                        if n>=len(a):
                            return
                        if k==a[n].pos1[2]:
                            poscur=[i,j,k]
                            a[n].val1= getval(pospre,poscur) #i,j,k ->i0,j0,k0
                            pospre=poscur
                            n=n+1
                else:
                    None
        else:
            None
## -- 2.2 search the .cube for the value of these numbers (tested) --

# ---- 2 search for the value of a point in the .cube ----


# --- 3 all the points passed by a plan defined by three points ---


def dotproduct(X1,X2): #Majuscule signifies a vector in 3D
    """mathematic dot product for two vectors"""
    x=X1[1]*X2[2]-X1[2]*X2[1]
    y=X1[2]*X2[0]-X1[0]*X2[2]
    z=X1[0]*X2[1]-X1[1]*X2[0]
    return [x,y,z]

def zerotest(X): # to test if a vector is zero
    d=0
    for s in X:
        d=d+abs(s)
    if 0<=d<1e-10:
        return 0
    else:
        return 1

class Ptvalue:
    "associate a value to a position 3D in the molecule space"
    def __init__(self):
        self.pos=[0]*3
        self.pos1=[0]*3
        self.pos2=[0]*3
        self.r1=0
        self.r2=0
        self.val1=0
        self.val2=0
        self.val=0
        # self.vlsort is just for the sort in part 2.1
        self.vlsort=0
        
    def affiche(self):
        print "pos:",self.pos,"vlsort:",self.vlsort,'pos1:',self.pos1,\
              'val1:',self.val1,'\n'
        
def VectorPer(X1,X2,X3):
    """search for the vertical vector (relative position) to the plan"""
    X1=np.array(X1)#convert the type to facilite the calcul
    X2=np.array(X2)
    X3=np.array(X3)
    orispace=np.array(mole.ori) #the origin of the space
    mole.pas=np.array(mole.pas)

    V1=(X1-orispace)/mole.pas #We use the space relative to the real one
    V2=(X2-orispace)/mole.pas
    V3=(X3-orispace)/mole.pas

    Vplan=dotproduct(V2-V1,V3-V1)
    if zerotest(Vplan)==0:
        Vplan=dotproduct(V1-V2,V3-V2) 
        return Vplan,V2
    else:
        return Vplan,V1

def Searchdir(X1,X2,X3):
    """search for the optimal directions in which it contains more points"""
    global drc,flag
    drc=range(3) #0,1,2 corresponds respectively the three directions x,y,z

    Vplan,Vori=VectorPer(X1,X2,X3)

    #search for the two directions that we want to fix
    maximum=max(Vplan)
    minimum=min(Vplan)
    if maximum>=0 and minimum>=0:  
        flag=Vplan.index(maximum)
    if maximum<=0 and minimum<=0:
        flag=Vplan.index(minimum)
    else:
        temp=maximum+minimum
        if temp>=0:
            flag=Vplan.index(maximum)
        else:
            flag=Vplan.index(minimum)
    del drc[flag] #'flag' indicate the direction that we don't need

    #drc is the list containing the two directions that we will work with

#Principal function in this part, we could find all the points concerned
#without its value in the optimal direction, and 'drc,flag,Ptvl' are global

def SearchPt(X1,X2,X3):
    """search for all the positions of the concerned point"""
    global Ptvl #Ptvl stocks the values of the point
    Vplan,Vori=VectorPer(X1,X2,X3)
    Searchdir(X1,X2,X3)
    Ptvl=[]

    #(Vpt-Vori).Vplan=0
    summe=0            # 
    for i in range(3):
        summe+=Vori[i]*Vplan[i]
    for i in range(mole.nbpas[drc[0]]):
        for j in range(mole.nbpas[drc[1]]):
            vect=[0]*3
            vect[drc[0]]=i
            vect[drc[1]]=j
            temp=summe-Vplan[drc[0]]*i-Vplan[drc[1]]*j
            vect[flag]=float(temp)/Vplan[flag]

            Ptvl0=Ptvalue()
            Ptvl0.pos=vect

            Ptvl.append(Ptvl0)
 

# ---- 3 all the concerned points have been found (simply tested) ----

# --- 4 draw the contours ---
import matplotlib
import matplotlib.cm as cm
import matplotlib.mlab as mlab
import matplotlib.pyplot as plt

Nbcontour= 15 #This is the number of the contours that we want to draw
Zoom=10e8

def trace():
    """this is the command directly used by the button"""
    temp=[]
    for i in range(len(Ptvl)):
        temp.append(Ptvl[i].val1)
    maximum=max(temp)*Zoom
    minimum=min(temp)*Zoom

    V=np.arange(minimum,maximum,(maximum-minimum)/Nbcontour)

    nspx=mole.nbpas[drc[0]]
    nspy=mole.nbpas[drc[1]]
    
    temp=((Ptvl[0].pos1[flag]-Ptvl[100].pos1[flag])*mole.pas[flag])**2
    dspx=math.sqrt(((Ptvl[0].pos1[drc[0]]-Ptvl[100].pos1[drc[0]])\
                    *mole.pas[drc[0]])**2+temp)
    dspx=dspx/abs(Ptvl[0].pos1[drc[0]]-Ptvl[100].pos1[drc[0]])
    dspy=math.sqrt(((Ptvl[0].pos2[drc[0]]-Ptvl[100].pos2[drc[0]])\
                    *mole.pas[drc[0]])**2+temp)
    dspy=dspy/abs(Ptvl[0].pos2[drc[0]]-Ptvl[100].pos2[drc[0]])

    xori,yori=0,0 #could be changed

    x=np.arange(xori*dspx,(xori+nspx)*dspx,dspx) 
    y=np.arange(yori*dspy,(yori+nspy)*dspy,dspy)


    Z=numpy.matlib.zeros((mole.nbpas[drc[1]],mole.nbpas[drc[0]])\
                         ,dtype=float)
    for i in range(len(Ptvl)):
        Z[Ptvl[i].pos1[drc[1]],Ptvl[i].pos1[drc[0]]]=Ptvl[i].val1*Zoom

    plt.figure()
    CS = plt.contour(x,y,Z,V,linewidths=1)
    plt.clabel(CS,fontsize=10)
    plt.title('Density of spin for molecule (Scale  1:'+str(Zoom)+')')

    plt.show()


# ---- 4 
            
        
def test():
    X2=np.array([0,0,1])+np.array(mole.ori)
    X1=mole.ori
    X3=np.array([0,1,0])+np.array(mole.ori)
    SearchPt(X1,X2,X3)
    Ptcube(Ptvl)
    #qsort(Ptvl,2)
    searchval1(Ptvl)
    print 'The information has been successfully recorded!\n'
    fI.close()

    
# ------ Principal programme ------

if __name__ == '__main__':
    fen = Tk()
    fen.title("Calcul of the spin density")
    
    Button(fen,relief='raised',border=2,bg='lightgreen',text = "1.Choose the file .cube of the molecule", \
           command=demo_parcours).pack(side='top',fill=X)

    Button(fen,relief='raised',border=2,bg='lightgreen',text = "2.Register the information of molecule", \
           command=moleinfo).pack(side='top',fill=X)

    Button(fen,relief='raised',border=2,bg='lightgreen',text = "3.Register the plane information", \
           command=test).pack(side='top',fill=X)

    Button(fen,relief='raised',border=2,bg='lightgreen',text = "4.draw the contours", command = trace
           ).pack(side='top',fill=X)

    Button(fen,relief='ridge',border=3,fg='red',text='Quit',command=fen.destroy\
           ).pack(side='top',fill=X)
    
    tkMessageBox.showinfo("Brief Introduction","     This executive program works for\nthe visualization of the spin density\nin a specific section\n     To begin with, please choose the\naddress of the *.cube file")

    fen.mainloop()
    

